"""
The main script that contains subclass of cplex class.
All customized featured are to be added incrementally

Created: May 1st?

Version 1.0 May 7th

Author Derek Zhang
"""
import cplex
import random
from math import floor
from heapq import nlargest

class fscplex(cplex.Cplex):
    def __init__(self):
        cplex.Cplex.__init__(self) #has to be called to activate the constructor of Cplex object
        self.NI = 1
        self.NS = 1
        self.eps = 0.1
        self.I = range(self.NI)
        self.J = range(self.NS)
        self.x_name = []
        self.v_name = []
        self.y_name = []
        self.dm_name = []
        self.vb_name = []
        self.rhs = []
        
        self.num_bigm = 0
        self.bigm_rowname = []
        self.bigm_colname = []
        self.bigm_coef = []
        self.bigm_ind = []

        self.v_lb = []
        self.v_hlb = []

        self.p = 0
        self.xi = []
        self.xi_ind = []

        self.xi_max=[]
        self.xi_min=[]
        self.x_heur=[]

        self.flag = [1]
        self.incumbent = [1]
        
    def register_callback(self, mycallback, Print=False):
        cbinstance = cplex.Cplex.register_callback(self, mycallback)
        cbinstance.times_called = 0
        cbinstance.Print = Print
        cbinstance.x_name = self.x_name
        cbinstance.v_name = self.v_name
        cbinstance.y_name = self.y_name
        cbinstance.vb_name = self.vb_name
        cbinstance.v_lb = self.v_lb
        cbinstance.v_hlb = self.v_hlb
        cbinstance.flag = self.flag
        cbinstance.incumbent = self.incumbent
        cbinstance.NI = self.NI
        cbinstance.NS = self.NS
        cbinstance.I = self.I
        cbinstance.S = self.S
        cbinstance.p = self.p
        cbinstance.rhs = self.rhs
        cbinstance.xi = self.xi
        cbinstance.xi_max = self.xi_max
        cbinstance.xi_min = self.xi_min
        cbinstance.dm_name = self.dm_name
        cbinstance.bigm_rowname = self.bigm_rowname
        cbinstance.bigm_colname = self.bigm_colname
        cbinstance.bigm_coef = self.bigm_coef
        cbinstance.bigm_ind = self.bigm_ind
        cbinstance.num_bigm = self.num_bigm
        cbinstance.node_x_val = [0] * self.NI
        cbinstance.node_y_val = [0] * self.NS
        cbinstance.slack = [0] * self.num_bigm

        return cbinstance

    def generate_instance(self, NI=5, NS=5, eps=0.3, Ordered=False):
        random.seed(100)
        self.NI = NI
        self.NS = NS
        self.eps = eps
        self.I = range(NI)
        self.S = range(NS)
        self.x_name = ["x" + '_' + str(i) for i in self.I]
        self.v_name = ["v" + '_' + str(i) for i in self.I]
        self.y_name = ["y" + '_' + str(s) for s in self.S]
        self.v_hlb = [0] * NI
        self.p = int(floor(eps*NS)) #!!!!!! it is not int(floor(eps*NS)) big bug found on May 9th

        cost = [random.uniform(0,1) for i in self.I]
        prob = [1.0/NS]*NS

        # demand to be sorted later in terms of each continuous variable
        if Ordered:
            self.rhs = [[[] for i in self.I] for s in self.S]
            for i in self.I:
                i_rhs = [random.normalvariate(5,1) for s in self.S]
                i_rhs.sort(reverse=True)
                #i_rhs.sort()
                for s in self.S:
                    self.rhs[s][i] = i_rhs[s]
                # To tweak the priority list
                #self.rhs[0][i] = random.normalvariate(10,2)
        else:
            self.rhs = [[random.normalvariate(5,1) for i in self.I] for s in self.S]
            #self.rhs = [[random.uniform(4,6) for i in self.I] for s in self.S]

        demand = self.rhs
        
        #Create variables 
        x_name = self.x_name
        x_type = ["C"] * NI
        x_lb   = [0] * NI
        x_obj  = cost

        v_name = self.v_name
        v_type = ["C"] * NI
        v_lb   = [0] * NI
        v_obj =  [0] * NI

        y_name = self.y_name
        y_type = ["B"] * NS
        y_lb   = [0] * NS
        y_obj  = [0] * NS #used to be 0

        self.objective.set_sense(self.objective.sense.minimize)
        # Add columns in one pass
        self.variables.add(  obj = x_obj +  v_obj  + y_obj,
                        types = x_type + v_type + y_type,
                        names = x_name + v_name + y_name,
                           lb = x_lb + v_lb + y_lb)
        # Add rows
        # knapsack constraint
        expr = [[y_name, prob]]
        sen =  ["L"]
        rhs =  [eps]
        names = ["KP"]
        self.linear_constraints.add( lin_expr = expr,
                                  senses = sen,
                                  rhs = rhs,
                                  names = names)
        # capacity constraint
        self.dm_name = ["dm_" + str(i) + '_s_' +str(s) for s in self.S for i in self.I]

        expr = [ [ [v_name[i], y_name[s]], [1, demand[s][i]] ] for s in self.S for i in self.I]
        sen =  ["G"] * NI * NS
        rhs = [ demand[s][i] for s in self.S for i in self.I]
        names = self.dm_name
        self.linear_constraints.add( lin_expr = expr,
                                  senses = sen,
                                  rhs = rhs,
                                  names = names)
        # bounds constraint
        expr = [ [ [x_name[i], v_name[i]], [1, -1] ] for i in self.I]
        sen = ["G"] * NI
        rhs = [0] * NI
        names = ["vb"+str(i) for i in self.I]
        self.linear_constraints.add( lin_expr = expr,
                                  senses = sen,
                                  rhs = rhs,
                                  names = names)
        # set bigm names
        self.num_bigm = len(self.dm_name)
        self.bigm_ind = [0] * self.num_bigm
        self.bigm_rowname = self.dm_name
        self.bigm_colname = [ 'y_'+str(s) for s in self.S for i in self.I]
        self.bigm_coef = self.linear_constraints.get_coefficients(zip(self.bigm_rowname, self.bigm_colname))

        #self.write("ccfs.lp")

    def get_x_heur(self):
        # solve the problem directly by lp relaxation to obtain heurisitc solution
        lp = cplex.Cplex(self)
        lp.set_problem_type(lp.problem_type.LP)
        lp.solve()
        self.x_heur = lp.solution.get_values(self.x_name)
        lp = []

    def get_hard_v_lower_bounds(self):
        p = 0  #follow the notation on the paper
        xi_array = []
        for i in self.I:
            print "----\n"
            rowname = ["dm_" + str(i) + '_s_' +str(s) for s in self.S]
            xi_array = self.linear_constraints.get_rhs(rowname)

            xi_array.sort(reverse=True)
            print xi_array
            p = int(floor(self.eps * self.NS))
            self.v_hlb[i] = xi_array[p]

        print "\n\n***********"
        print self.v_hlb

    # Set initial bigm values
    def set_bigm(self, bigm_val):
        if len(bigm_val) != self.num_bigm:
            print "Wrong"
            return 
        self.linear_constraints.set_coefficients(zip(self.bigm_rowname, self.bigm_colname, bigm_val))

        self.bigm_coef = self.linear_constraints.get_coefficients(zip(self.bigm_rowname, self.bigm_colname))


    # Increase bigm values
    def increase_bigm(self, step_length=0.0, step_ratio=0.0, batch = False):

        #print step_length
        #print self.linear_constraints.get_coefficients(zip(self.bigm_rowname, self.bigm_colname))

        if batch == False:
            for i in range(self.num_bigm):
                bigm_val = self.linear_constraints.get_coefficients(self.bigm_rowname[i], self.bigm_colname[i])
                bigm_val_new = 0.0

                if self.bigm_ind[i] == 1:
                    #print i
                    local_rhs = self.linear_constraints.get_rhs(self.bigm_rowname[i])

                    if step_length != 0.0:
                        bigm_val_new = min(bigm_val + step_length, local_rhs)
                    elif step_ratio != 0.0:
                        bigm_val_new = min(bigm_val + step_ratio * local_rhs, local_rhs)
                        #print "bigm_val before modification ", bigm_val
                        #local_i = i%self.NI

                        ### We bound the maximum of bigm to be set by the corresponding rhs value (if not setting then the iteration will be too many!!!)

                    self.linear_constraints.set_coefficients(self.bigm_rowname[i], self.bigm_colname[i], bigm_val_new)

                        #print "bigm_val after modification", self.linear_constraints.get_coefficients(self.bigm_rowname[i], self.bigm_colname[i])
                    self.bigm_coef[i] = self.linear_constraints.get_coefficients(self.bigm_rowname[i], self.bigm_colname[i])                    

        elif batch == True:
            for i in range(self.NI):
                chg_ind = 0
                for s in self.S:
                    ind = i+ s*self.NI
                    if self.bigm_ind[ind] == 1:
                        chg_ind = 1
                        break
                    
                if chg_ind == 1:
                    for s in self.S:
                        ind = i+ s*self.NI
                    
                        bigm_val = self.linear_constraints.get_coefficients(self.bigm_rowname[ind], self.bigm_colname[ind])
                        bigm_val_new = 0.0

                        #if self.bigm_ind[ind] == 1:
                        if 1:
                            local_rhs = self.linear_constraints.get_rhs(self.bigm_rowname[ind])
                            if step_length != 0.0:
                                bigm_val_new = min(bigm_val + step_length, local_rhs)
                            elif step_ratio != 0.0:
                                bigm_val_new = min(bigm_val + step_ratio * local_rhs, local_rhs)
                            self.linear_constraints.set_coefficients(self.bigm_rowname[ind], self.bigm_colname[ind], bigm_val_new)
                            self.bigm_coef[ind] = self.linear_constraints.get_coefficients(self.bigm_rowname[ind], self.bigm_colname[ind])                    
                            
                
    def set_bigm_two(self, val1, val2):
        if self.NI != 2:
            return "Wrong"
        bigm_val1 = [val1] * self.NS
        bigm_val2 = [val2] * self.NS
        bigm_rowname1 = ['dm_0' + '_s_' +str(s) for s in self.S]
        bigm_colname1 = ['y_' + str(s) for s in self.S]
        bigm_rowname2 = ['dm_1' + '_s_' +str(s) for s in self.S]
        bigm_colname2 = ['y_' + str(s) for s in self.S]

        self.linear_constraints.set_coefficients(zip(bigm_rowname1, bigm_colname1, bigm_val1))
        self.linear_constraints.set_coefficients(zip(bigm_rowname2, bigm_colname2, bigm_val2))
        
    # find the pth largest element of xi
    def get_xi(self):
        self.xi = [0] * self.NI
        self.xi_max = [0] * self.NI
        self.xi_min = [0] * self.NI
        for i in range(self.NI):
            xi_rowname = ["dm_" + str(i) + '_s_' +str(s) for s in self.S]
            xi_colname = self.y_name
            xi_temp = self.linear_constraints.get_coefficients(zip(xi_rowname, xi_colname))
            xi_temp.sort(reverse=True)
            self.xi[i] = xi_temp[self.p]
            print self.xi[i]
            self.xi_max[i] = xi_temp[0]
            self.xi_min[i] = xi_temp[-1]

    # Strenghten the formulation 
    def strengthen_formulation(self):
        self.get_xi()
        bigm_val = [0] * self.num_bigm
        for s in self.S:
            for i in self.I:
                ind = i + s*self.NI
                if self.rhs[s][i] >= self.xi[i]:
                    bigm_val[ind] = self.rhs[s][i] - self.xi[i]
                elif self.rhs[s][i] < self.xi[i]:
                    bigm_val[ind] = 0
                    self.linear_constraints.set_rhs(self.bigm_rowname[ind], 0)
                    
        self.linear_constraints.set_coefficients(zip(self.bigm_rowname, self.bigm_colname, bigm_val))

        self.bigm_coef = self.linear_constraints.get_coefficients(zip(self.bigm_rowname, self.bigm_colname))
                




